package cn.swing.main.srv.cv.onnx;

import cn.swing.main.srv.cv.BaseOpencv;
import ai.onnxruntime.*;
import org.opencv.core.*;
import org.opencv.highgui.HighGui;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import java.nio.FloatBuffer;
import java.text.DecimalFormat;
import java.util.*;
import java.util.stream.Collectors;

/*
这是抽象出来的父类，不需要进行修改
*/
public class PtOnnx extends BaseOpencv {

    public static OrtEnvironment environment;

    public static OrtSession session;

    public static long count;

    public static long channels;

    public static long netHeight;

    public static long netWidth;

    public static float confThreshold = 0.35f;//置信度阈值

    public static float nmsThreshold = 0.45f;//交并比iou阈值

    static Mat img;

    public static float srcw;

    public static float srch;

    public static float scaleW;

    public static float scaleH;

    //提供获取置信度和交并比的函数
    public static float getConfThreshold() {
        return confThreshold;
    }

    public static void setConfThreshold(float confThreshold) {
        PtOnnx.confThreshold = confThreshold;
    }

    public static float getNmsThreshold() {
        return nmsThreshold;
    }

    public static void setNmsThreshold(float nmsThreshold) {
        PtOnnx.nmsThreshold = nmsThreshold;
    }

    public static void main(String[] args) throws Exception {
        String img = "D:\\swing\\label\\a8bd3ae4-3a97-4962-a9bf-67fc528f262c-1\\2309.jpg";
        String modelPath = "D:\\swing\\yolo11n.onnx";
        String labelStr = "tennis racket/person";
        final Map<Object, Object> objectObjectMap = pt_onnx(img, modelPath, labelStr);
    }

    /*
    该函数是主函数，将所有函数穿起来
     */
    public static Map<Object, Object> pt_onnx(String imgPath, String modelPath, String labelstr) throws Exception {
        //获取run（）的数据
        float[][] data = ONNX_load_run(imgPath, modelPath);
        //对数据进行塑形
        float[][] rawdata = shape_rebuild(data);
        //模型后处理阶段
        return photo_process(rawdata, imgPath, labelstr);

    }

    /*
     * 模型初始化及模型调用获取数据
     * */
    public static float[][] ONNX_load_run(String imgPath, String modelPath) throws Exception {
        //读取模型和对话类
        load(modelPath);
        //用于读取图像文件并加载到内存中
        img = Imgcodecs.imread(imgPath);
        //获取图片的长宽和及其和模型的比例
        srcw = img.width();
        srch = img.height();
        scaleW = srcw / netWidth;
        scaleH = srch / netHeight;

        // mat是opencv库中常用对象，表示图像视频矩阵
        // 该函数用户获取图片路径并且将其转换为矩阵形式
        Mat dst = new Mat();
        // 这段代码的作用是将输入图像 img 调整为指定的尺寸，并将结果保存在 dst 中。
        Imgproc.resize(img, dst, new Size(netWidth, netHeight));
        // 转换成Tensor数据格式
        OnnxTensor tensor = transferTensor(dst);
        OrtSession.Result result = session.run(Collections.singletonMap("images", tensor));
        OnnxTensor res = (OnnxTensor) result.get(0);
        float[][][] dataRes = (float[][][]) res.getValue();
        float[][] data = dataRes[0];
        return data;
    }

    /*
     * shape整形，每种模型不一样，需要根据模型特点来进行重写
     * */
    public static float[][] shape_rebuild(float[][] data) {
        float rawData[][] = new float[data.length][6];
        for (int i = 0; i < data.length; i++) {
            for (int j = 0; j < 5; j++) {
                rawData[i][j] = data[i][j];
            }
            float[] conditionalProbabilities = Arrays.copyOfRange(data[i], 5, data[0].length);
            int label = argmax(conditionalProbabilities);
            rawData[i][5] = label;
        }
        return rawData;
    }

    /*
     * 该函数是模型后处理函数，封装了置信度处理，交并比处理，裱框处理
     * 其输入应是一个[x1,y1,x2,y2,conf,label(害虫的数组下标)]标准类型
     * 输出一个map，输入是经过塑性后的data，图片原地址用于保存，label用于获取标签
     *
     * 可泛化
     * */
    public static Map<Object, Object> photo_process(float[][] data, String imgPath, String labelstr) throws Exception {


         /*
        (1)将通过模型识别后产生的数据进行第一轮置信度筛选，
        将选出的结果存到哈希表中：
        key为label的值，List为每一个框的参数
         */

        Map<Integer, List<float[]>> class2Bbox = new HashMap<>();
        for (float[] bbox : data) {
            //获取置信度，如果小于所设值则从头开始循环
            float score = bbox[4];
            if (score < confThreshold) {
                continue;
            }
            // 截取数组的后面三个值，这三个值代表了该框与模型中的哪种目标相似度最高
            //获取相似度最高值的数组下标
            int label = (int) bbox[5];
            // xywh to (x1, y1, x2, y2)
            xywh2xyxy(bbox);
            // 如果左（下）边框值大于右（上）边框，则是错误裱框
            if (bbox[0] >= bbox[2] || bbox[1] >= bbox[3]) {
                continue;
            }
            //将label的值作为键，合格的框加入到该label下对应的值中
            class2Bbox.putIfAbsent(label, new ArrayList<>());
            class2Bbox.get(label).add(bbox);
        }

        /*
         * （2）这一阶段进行非极大值抑制，最终得到一个list数组包含
         * 修改了detection的tostring函数输出的结构为[x1,y1,x2,y2,conf,label]
         *
         * */
        //获取英文字符串并分割变成数组
        String[] labels = labelstr.split("/");

        List<Detection> detections = new ArrayList<>();
        //遍历class2Bbox中所有的键
        for (Map.Entry<Integer, List<float[]>> entry : class2Bbox.entrySet()) {
            List<float[]> bboxes = entry.getValue();
            //非极大值抑制，nmsThreshold代表交并比iou，删除相似裱框只留一个置信度最高的
            bboxes = nonMaxSuppression(bboxes, nmsThreshold);
            for (float[] bbox : bboxes) {
//                String labelString = labels[entry.getKey()];
                detections.add(new Detection(String.valueOf(entry.getKey()), entry.getKey(),
                        Arrays.copyOfRange(bbox, 0, 4), bbox[4]));
            }
        }
        /*
         * （3）这一阶段进行图片裱框
         * */
        Imgproc.resize(img, img, new Size(srcw, srch));
        for (Detection detection : detections) {
            float[] bbox = detection.getBbox();
            float x1 = bbox[0] * scaleW;
            float y1 = bbox[1] * scaleH;
            float x2 = bbox[2] * scaleW;
            float y2 = bbox[3] * scaleH;
            float config = detection.getConfidence();
            String className = detection.getLabel();
            Point point1 = new Point(x1, y1);
            Point point2 = new Point(x2, y2);
            Imgproc.rectangle(img, point1, point2, new Scalar(0, 0, 255), 2);
            String conf = new DecimalFormat("#.###").format(config);
            Imgproc.putText(img, className + " " + conf, new Point(x1, y1 - 5), 0, 0.5, new Scalar(255, 0, 0), 1);
        }
        HighGui.imshow("Frame", img);
        HighGui.waitKey(0);
        // 保存图像到同级目录
//        Imgcodecs.imwrite(imgPath, img);
        // 将检测的数组放入map中传到下一个函数
        Map<Object, Object> map = new HashMap<Object, Object>();
        map.put("boxes", detections);
        return map;

    }

    /*
     * getsum函数也需要泛化
     * */
    //    public static String getSum(Map<Object, Object> map, JSONObject names) throws JsonProcessingException {
    //        List<Detection> boxes = (List<Detection>) map.get("boxes");
    //        //将每个列表第2元素取出变为string
    //        List<Integer> listtest = new ArrayList<>();
    //        for (Detection list : boxes) {
    //            listtest.add(list.getClsId());
    //        }
    //        Integer[] integersArray = listtest.toArray(new Integer[0]);
    //        //获取键对应值的数组
    //        String[] Array = new String[integersArray.length];
    //        for (int i = 0; i < integersArray.length; i++) {
    //            // 根据第一个 JSON 对象的键找到对应的值
    //            String value = (String) names.get(String.valueOf(integersArray[i]));
    //            // 将值添加到新数组中
    //            Array[i] = value;
    //        }
    //        Map<String, Integer> countMap = new HashMap<>();
    //        Map<String, String> finalMap = new HashMap<>();
    //        // 遍历数组，统计元素出现的次数
    //        for (String num : Array) {
    //            if (countMap.containsKey(num)) {
    //                // 如果 Map 中已存在该元素，则将其次数加一
    //                countMap.put(num, countMap.get(num) + 1);
    //            } else {
    //                // 如果 Map 中不存在该元素，则将其次数设为 1
    //                countMap.put(num, 1);
    //            }
    //        }
    //        // 打印结果
    //        for (Map.Entry<String, Integer> entry : countMap.entrySet()) {
    //            //            System.out.println("元素 " + entry.getKey() + " 出现的次数为：" + entry.getValue());
    //            finalMap.put(entry.getKey(), String.valueOf(entry.getValue()));
    //        }
    //
    //        // 创建 ObjectMapper 对象
    //        ObjectMapper objectMapper = new ObjectMapper();
    //
    //        // 将 Map 转换为 JSON 字符串
    //        String jsonString = objectMapper.writeValueAsString(finalMap);
    //        System.out.println(jsonString);
    //        return jsonString;
    //    }

    //    /*将一个以”/“分隔的字符串转换成json格式用于后续输出
    //    这个函数仅用来获取中文字符串的json格式，可以进行泛化
    //     */
    //    public static JSONObject array_transform(String str) {
    //        //先将其变成数组
    //        String[] array = str.split("/");
    //        // 创建一个 JSON 对象
    //        JSONObject jsonObject = new JSONObject();
    //
    //        // 遍历数组，将每个元素添加到 JSON 对象中
    //        for (int i = 0; i < array.length; i++) {
    //            jsonObject.put(String.valueOf(i), array[i]);
    //        }
    //        // 输出 JSON 对象
    //        //        System.out.println(jsonObject.toJSONString());
    //        JSONObject json = JSONObject.parseObject(jsonObject.toJSONString());
    //        return json;
    //
    //    }

    public static void xywh2xyxy(float[] bbox) {
        float x = bbox[0];
        float y = bbox[1];
        float w = bbox[2];
        float h = bbox[3];

        bbox[0] = x - w * 0.5f;
        bbox[1] = y - h * 0.5f;
        bbox[2] = x + w * 0.5f;
        bbox[3] = y + h * 0.5f;
    }

    public static List<float[]> nonMaxSuppression(List<float[]> bboxes, float iouThreshold) {
        List<float[]> bestBboxes = new ArrayList<>();
        bboxes.sort(Comparator.comparing(a -> a[4]));
        while (!bboxes.isEmpty()) {
            float[] bestBbox = bboxes.remove(bboxes.size() - 1);
            bestBboxes.add(bestBbox);
            bboxes = bboxes.stream().filter(a -> computeIOU(a, bestBbox) < iouThreshold).collect(Collectors.toList());
        }
        return bestBboxes;
    }

    public static float computeIOU(float[] box1, float[] box2) {

        float area1 = (box1[2] - box1[0]) * (box1[3] - box1[1]);
        float area2 = (box2[2] - box2[0]) * (box2[3] - box2[1]);

        float left = Math.max(box1[0], box2[0]);
        float top = Math.max(box1[1], box2[1]);
        float right = Math.min(box1[2], box2[2]);
        float bottom = Math.min(box1[3], box2[3]);

        float interArea = Math.max(right - left, 0) * Math.max(bottom - top, 0);
        float unionArea = area1 + area2 - interArea;
        return Math.max(interArea / unionArea, 1e-8f);

    }

    //返回最大值的索引
    public static int argmax(float[] a) {
        float re = -Float.MAX_VALUE;
        int arg = -1;
        for (int i = 0; i < a.length; i++) {
            if (a[i] >= re) {
                re = a[i];
                arg = i;
            }
        }
        return arg;
    }

    public static void load(String model_path) throws Exception {
        // 加载ONNX模型环境类和会话类
        environment = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions sessionOptions = new OrtSession.SessionOptions();
        session = environment.createSession(model_path, sessionOptions);

        //输出张量相关信息
        Map<String, NodeInfo> infoMap = session.getInputInfo();
        TensorInfo nodeInfo = (TensorInfo) infoMap.get("images").getInfo();
        count = nodeInfo.getShape()[0];//1 模型每次处理一张图片
        channels = nodeInfo.getShape()[1];//3 模型通道数
        netHeight = nodeInfo.getShape()[2];//640 模型高
        netWidth = nodeInfo.getShape()[3];//640 模型宽
        //加载opencvjava包
        try {
            //运用于windows系统
            System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
            //运用于linux服务器中
            //            System.load("/usr/lib/libopencv_java490.so");
        } catch (Exception e) {
            System.out.println("load error");
        }

    }

    public static OnnxTensor transferTensor(Mat dst) {
        Imgproc.cvtColor(dst, dst, Imgproc.COLOR_BGR2RGB);
        dst.convertTo(dst, CvType.CV_32FC1, 1. / 255);
        float[] whc =
                new float[Long.valueOf(channels).intValue() * Long.valueOf(netWidth).intValue() * Long.valueOf(netHeight).intValue()];
        dst.get(0, 0, whc);
        float[] chw = whc2cwh(whc);
        OnnxTensor tensor = null;
        try {
            tensor = OnnxTensor.createTensor(environment, FloatBuffer.wrap(chw), new long[] {count, channels,
                    netWidth, netHeight});
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(0);
        }
        return tensor;
    }

    public static float[] whc2cwh(float[] src) {
        float[] chw = new float[src.length];
        int j = 0;
        for (int ch = 0; ch < 3; ++ch) {
            for (int i = ch; i < src.length; i += 3) {
                chw[j] = src[i];
                j++;
            }
        }
        return chw;
    }

}


